"""
# -*- coding: utf-8 -*-
# @Time    : 2023/10/17 8:37
# @Author  : 王摇摆
# @FileName: gnn.py
# @Software: PyCharm
# @Blog    ：https://blog.csdn.net/weixin_44943389?type=blog
"""

import dgl
import torch
import torch.nn as nn
import torch.optim as optim
from dgl.data import citation_graph

# 加载一个示例数据集（这里用的是Cora数据集）
data = citation_graph.load_cora()

# 创建图
g = data.graph
features = torch.FloatTensor(data.features)
labels = torch.LongTensor(data.labels)
train_mask = torch.BoolTensor(data.train_mask)
val_mask = torch.BoolTensor(data.val_mask)
test_mask = torch.BoolTensor(data.test_mask)

# 定义一个简单的GNN模型
class GCN(nn.Module):
    def __init__(self, in_feats, hidden_size, num_classes):
        super(GCN, self).__init__()
        self.conv1 = nn.GraphConv(in_feats, hidden_size)
        self.conv2 = nn.GraphConv(hidden_size, num_classes)

    def forward(self, g, features):
        x = torch.relu(self.conv1(g, features))
        x = self.conv2(g, x)
        return x

# 初始化模型、损失函数和优化器
model = GCN(in_feats=features.shape[1], hidden_size=16, num_classes=7)
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=0.01)

# 训练模型
def train():
    model.train()
    optimizer.zero_grad()
    logits = model(g, features)
    loss = criterion(logits[train_mask], labels[train_mask])
    loss.backward()
    optimizer.step()
    return loss.item()

# 验证模型
def evaluate():
    model.eval()
    with torch.no_grad():
        logits = model(g, features)
        val_loss = criterion(logits[val_mask], labels[val_mask])
        val_acc = (logits[val_mask].argmax(1) == labels[val_mask]).float().mean().item()
        return val_loss.item(), val_acc

# 训练模型并验证
for epoch in range(50):
    loss = train()
    val_loss, val_acc = evaluate()
    print(f'Epoch {epoch + 1}, Loss: {loss:.4f}, Val Loss: {val_loss:.4f}, Val Acc: {val_acc:.4f}')

# 在测试集上评估模型
model.eval()
with torch.no_grad():
    logits = model(g, features)
    test_loss = criterion(logits[test_mask], labels[test_mask])
    test_acc = (logits[test_mask].argmax(1) == labels[test_mask]).float().mean().item()
    print(f'Test Loss: {test_loss:.4f}, Test Acc: {test_acc:.4f}')
